package com.cyf.common.core.coll;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

/**
 * 带有最大容量限制的优先队列，如果队列已满，并且新元素小于队列中的最大元素，则删除最大元素并添加新元素。
 * @author chenyifan
 * @create 2024-09-23 13:43
 */
public class MinNHeap<E> {

    private PriorityQueue<E> queue;
    private Comparator<? super E> comparator;
    private int maxSize;

    public MinNHeap(int maxSize, Comparator<? super E> comparator) {
        this.maxSize = maxSize;
        this.comparator = comparator;
        // 大顶堆
        this.queue = new PriorityQueue<>(maxSize, comparator.reversed());
    }

    public void add(E e) {
       if (queue.size() < maxSize) {
           queue.offer(e);
       } else if (compare(queue.peek(), e)) {
           // 如果堆顶元素 > 新数，则删除堆顶，加入新数入堆
           queue.poll();
           queue.offer(e);
       }
    }

    public List<E> getList() {
        return new ArrayList<>(queue);
    }

    /**
     * 比较 e1, e2 大小
     * @param e1
     * @param e2
     * @return e1 > e2 返回 true，否则返回 false
     */
    private boolean compare(E e1, E e2) {
        if (comparator != null) {
            return comparator.compare(e1, e2) > 0;
        } else {
            return ((Comparable<? super E>) e1).compareTo(e2) > 0;
        }
    }
}
